/* Copyright (C) 2015-2016 HangZhou Zenzet Technology Co., Ltd.
 * All right reserved

 * File:crypto/locketrsa.c
 * Author:guojianchuan/max
 * Date:2016-06-02

 */

/* system header */
#include <stdio.h>
#include <string.h>

/* 3rd project header */
#include "openssl/evp.h"
#include "openssl/rsa.h"
#include "openssl/x509.h"
#include "openssl/bio.h"
#include "openssl/pem.h"

/* my project header */
#include "locketdef.h"
#include "locketerrno.h"
#include "locketlog.h"
#include "locketcommon.h"
#include "locketconvert.h"

#include "locketrsa.h"

/**
* @Function:LOCKET_CIPHER_CheckRSAParam
* @Author: guojianchuan/max
* @Date: 2016-06-27
* @Description: Check RSA param
* @caution: 
*/
static int LOCKET_CIPHER_CheckRSAParam(IN CIPHER_RSA_S *pstRSA)
{
    if (NULL == pstRSA)
    {
        return ERR_COMMON_INVALID_PARAM;
    }

    if ((NULL == pstRSA->pucKey) ||
            (NULL == pstRSA->pucInput) ||
            (NULL == pstRSA->pucOutput) ||
            (0 == pstRSA->iInputLen))
    {
        return ERR_COMMON_INVALID_PARAM;
    }

    if (ERR_COMMON_SUCCESS != LOCKET_CIPHER_CheckEncode (pstRSA->iInputEncode))
    {
        return ERR_COMMON_INVALID_PARAM;
    }

    if (ERR_COMMON_SUCCESS != LOCKET_CIPHER_CheckEncode (pstRSA->iOutputEncode))
    {
        return ERR_COMMON_INVALID_PARAM;
    }

    return ERR_COMMON_SUCCESS;
}

/**
* @Function:LOCKET_CIPHER_RSAEncrypt
* @Author: guojianchuan/max
* @Date: 2016-06-02
* @Description: RSA encrypt with public key  (PADDING: PKCS-OAEP)
* @caution:
*/
int LOCKET_CIPHER_RSAEncrypt (INOUT CIPHER_RSA_S *pstRSA)
{
    int iErr = ERR_COMMON_SUCCESS;

    /* Check param */
    iErr = LOCKET_CIPHER_CheckRSAParam (pstRSA);
    if (ERR_COMMON_SUCCESS != iErr)
    {
        return iErr;
    }

    /* Read RSA Public-Key */
    BIO *pstBio                     = NULL;
    RSA *pstRsa                     = NULL;
    do
    {
        pstBio = BIO_new_mem_buf (pstRSA->pucKey, strlen((char*)pstRSA->pucKey));
        if (NULL == pstBio)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("PEM_read_bio_RSA_PUBKEY failed. %s", LOCKET_ERR_GetString());
            break;
        }

		pstRsa = PEM_read_bio_RSA_PUBKEY(pstBio, NULL, NULL, NULL); /* 不能用PEM_read_bio_RSAPublicKey，原因未知 */
        if (NULL == pstRsa)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("PEM_read_bio_RSA_PUBKEY failed. %s", LOCKET_ERR_GetString());
            break;
        }
    } while (0);

    /* don't forget cleanup */
    BIO_free_all(pstBio);

    if (ERR_COMMON_SUCCESS != iErr)
    {
        return iErr;
    }

    /* Encrypt with RSA */
    EVP_PKEY_CTX *pstCtx             = NULL;
    EVP_PKEY *pstKey                 = NULL;
    unsigned char *pucConvertOutput  = NULL;
    unsigned char *pucConvertInput   = NULL;
    unsigned char *pRet              = NULL;
    int iOutputLen                   = 0;
    do
    {
        /* Convert input */
        int iConvertInputLen  = 0;
        if (ENCODE_BINARY != pstRSA->iInputEncode)
        {
            pucConvertInput= (unsigned char*) LOCKET_malloc (pstRSA->iInputLen);
            if (NULL == pucConvertInput)
            {
                iErr = ERR_COMMON_NO_MEM;
                break;
            }
            memset (pucConvertInput, 0, pstRSA->iInputLen);
            pRet = LOCKET_CIPHER_OtherToByte (pstRSA->iInputEncode, pstRSA->pucInput,
                                              pstRSA->iInputLen, pucConvertInput, &iConvertInputLen);
            if (NULL == pRet)
            {
                iErr = ERR_COMMON_CONVERT_FAILED;
                break;
            }
        }
        else
        {
            pucConvertInput  = pstRSA->pucInput;
            iConvertInputLen = pstRSA->iInputLen;
        }

        pstKey = EVP_PKEY_new();
        if(NULL == pstKey)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_new failed. %s", LOCKET_ERR_GetString());
            break;
        }

        /* 后面只需要free pstKey即可，pstRsa会随之被free掉  */
        iErr = EVP_PKEY_assign_RSA(pstKey, pstRsa);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_assign_RSA failed. %s", LOCKET_ERR_GetString());
            break;
        }

        pstCtx = EVP_PKEY_CTX_new(pstKey, NULL);
        if(NULL == pstCtx)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_CTX_new failed. %s", LOCKET_ERR_GetString());
            break;
        }

        iErr = EVP_PKEY_encrypt_init (pstCtx);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_encrypt_init failed. %s", LOCKET_ERR_GetString());
            break;
        }

        /*
           The macro EVP_PKEY_CTX_set_rsa_padding() sets the RSA padding mode for ctx.
           The pad parameter can take the value:
           RSA_PKCS1_PADDING for PKCS#1 padding,
           RSA_SSLV23_PADDING for SSLv23 padding,
           RSA_NO_PADDING for no padding,
           RSA_PKCS1_OAEP_PADDING for OAEP padding (encrypt and decrypt only),
           RSA_X931_PADDING for X9.31 padding (signature operations only) and RSA_PKCS1_PSS_PADDING (sign and verify only).
        */
        /* OEAP-Padding, 对输入的字节长度有要求 */
		/* pkcs1 输入长度要小于RSA_size(rsa) - 11;oaep 输入长度要小于RSA_size(rsa) - 41*/
        iErr = EVP_PKEY_CTX_set_rsa_padding(pstCtx, RSA_PKCS1_OAEP_PADDING);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_CTX_set_rsa_padding failed. %s", LOCKET_ERR_GetString());
            break;
        }

#if 0
        /* 这个接口调用过之后，后面会挂掉，原因未知，暂时去除 */
        /* 检测输出长度大小, 第二参数设置成NULL */
        iErr = EVP_PKEY_encrypt(pstCtx, NULL, (size_t*) &iOutputLen,  pucConvertInput, iConvertInputLen);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_encrypt failed. %s", LOCKET_ERR_GetString());
            break;
        }
#endif

        if (ENCODE_BINARY != pstRSA->iOutputEncode)
        {
            pucConvertOutput = (unsigned char*) LOCKET_malloc (iOutputLen);
            if (NULL == pucConvertOutput)
            {
                iErr = ERR_COMMON_NO_MEM;
                break;
            }
            memset (pucConvertOutput, 0, iOutputLen);
        }
        else
        {
            pucConvertOutput = pstRSA->pucOutput;
        }

        /* Openssl文档中指出，需要传入outbuffer的长度, 但是从实际测试情况来看，即使传入长度为0，关系也不大 */
        /* 因为上面的那个接口会段错误，所以暂时用RSA_size(pstRsa)替代 */
		iOutputLen = RSA_size(pstRsa); 
        iErr = EVP_PKEY_encrypt(pstCtx, pucConvertOutput,  (size_t*) &iOutputLen, pucConvertInput, iConvertInputLen);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_encrypt failed. %s", LOCKET_ERR_GetString());
            break;
        }

        if (pucConvertOutput != pstRSA->pucOutput)
        {
            pRet = LOCKET_CIPHER_ByteToOther (pstRSA->iOutputEncode, pucConvertOutput, 
                                              iOutputLen, pstRSA->pucOutput, &pstRSA->iOutputLen);
            if (NULL == pRet)
            {
                iErr = ERR_COMMON_CONVERT_FAILED;
                break;
            }
        }
        else
        {
            pstRSA->iOutputLen = iOutputLen;
        }
    } while (0);

    /* don't forget cleanup */
    if (ENCODE_BINARY != pstRSA->iInputEncode)
    {
        LOCKET_free ((void**) &pucConvertInput);
    }
    if (ENCODE_BINARY != pstRSA->iOutputEncode)
    {
        LOCKET_free ((void**) &pucConvertOutput);
    }
    EVP_PKEY_CTX_free(pstCtx);
    EVP_PKEY_free(pstKey);

    return iErr;
}

/**
* @Function:LOCKET_CIPHER_RSADecrypt
* @Author: guojianchuan/max
* @Date: 2016-06-07
* @Description: RSA decrypt with private key (PADDING: PKCS-OAEP)
* @caution:
*/
int LOCKET_CIPHER_RSADecrypt (INOUT CIPHER_RSA_S *pstRSA)
{
    int iErr = ERR_COMMON_SUCCESS;

    /* Check param */
    iErr = LOCKET_CIPHER_CheckRSAParam (pstRSA);
    if (ERR_COMMON_SUCCESS != iErr)
    {
        return iErr;
    }

    /* Read RSA Private-Key */
    BIO *pstBio             = NULL;
    RSA *pstRsa             = NULL;
    do
    {
        pstBio = BIO_new_mem_buf (pstRSA->pucKey, strlen((char*) pstRSA->pucKey));
        if (NULL == pstBio)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("PEM_read_bio_RSA_PUBKEY failed. %s", LOCKET_ERR_GetString());
            break;
        }

		pstRsa = PEM_read_bio_RSAPrivateKey(pstBio, NULL, NULL, NULL);
        if (NULL == pstBio)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("PEM_read_bio_RSA_PUBKEY failed. %s", LOCKET_ERR_GetString());
            break;
        }
    } while (0);

    /* don't forget cleanup */
    BIO_free_all(pstBio);

    if (ERR_COMMON_SUCCESS != iErr)
    {
        return iErr;
    }

    /* Encrypt with RSA */
    EVP_PKEY_CTX *pstCtx             = NULL;
    EVP_PKEY *pstKey                 = NULL;
    unsigned char *pucConvertInput   = NULL;
    unsigned char *pucConvertOutput  = NULL;
    unsigned char *pRet              = NULL;
    int iOutputLen                   = 0;
    do
    {
        /* Convert input */
        int iConvertInputLen  = 0;
        if (ENCODE_BINARY != pstRSA->iInputEncode)
        {
            pucConvertInput= (unsigned char*) LOCKET_malloc (pstRSA->iInputLen);
            if (NULL == pucConvertInput)
            {
                iErr = ERR_COMMON_NO_MEM;
                break;
            }
            memset (pucConvertInput, 0, pstRSA->iInputLen);
            pRet = LOCKET_CIPHER_OtherToByte (pstRSA->iInputEncode, pstRSA->pucInput,
                                              pstRSA->iInputLen, pucConvertInput, &iConvertInputLen);
            if (NULL == pRet)
            {
                iErr = ERR_COMMON_CONVERT_FAILED;
                break;
            }
        }
        else
        {
            pucConvertInput  = pstRSA->pucInput;
            iConvertInputLen = pstRSA->iInputLen;
        }

        pstKey = EVP_PKEY_new();
        if(NULL == pstKey)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_new failed. %s", LOCKET_ERR_GetString());
            break;
        }

        iErr = EVP_PKEY_assign_RSA(pstKey, pstRsa);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_assign_RSA failed. %s", LOCKET_ERR_GetString());
            break;
        }

        pstCtx = EVP_PKEY_CTX_new(pstKey, NULL);
        if(NULL == pstCtx)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_CTX_new failed. %s", LOCKET_ERR_GetString());
            break;
        }

        iErr = EVP_PKEY_decrypt_init (pstCtx);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_decrypt_init failed. %s", LOCKET_ERR_GetString());
            break;
        }

        /*
           The macro EVP_PKEY_CTX_set_rsa_padding() sets the RSA padding mode for ctx.
           The pad parameter can take the value:
           RSA_PKCS1_PADDING for PKCS#1 padding,
           RSA_SSLV23_PADDING for SSLv23 padding,
           RSA_NO_PADDING for no padding,
           RSA_PKCS1_OAEP_PADDING for OAEP padding (encrypt and decrypt only),
           RSA_X931_PADDING for X9.31 padding (signature operations only) and RSA_PKCS1_PSS_PADDING (sign and verify only).
        */
        iErr = EVP_PKEY_CTX_set_rsa_padding(pstCtx, RSA_PKCS1_OAEP_PADDING);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_CTX_set_rsa_padding failed. %s", LOCKET_ERR_GetString());
            break;
        }

        /* 检测输出长度大小, 第二参数设置成NULL */ /* 这个地方不会挂掉 */
        iErr = EVP_PKEY_decrypt(pstCtx, NULL, (size_t*) &iOutputLen,  pucConvertInput, iConvertInputLen);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_decrypt failed. %s", LOCKET_ERR_GetString());
            break;
        }

        if (ENCODE_BINARY != pstRSA->iOutputEncode)
        {
            pucConvertOutput = (unsigned char*) LOCKET_malloc (iOutputLen);
            if (NULL == pucConvertOutput)
            {
                iErr = ERR_COMMON_NO_MEM;
                break;
            }
            memset (pucConvertOutput, 0, iOutputLen);
        }
        else
        {
            pucConvertOutput = pstRSA->pucOutput;
        }

        iErr = EVP_PKEY_decrypt(pstCtx, pucConvertOutput, (size_t*) &iOutputLen, pucConvertInput, iConvertInputLen);
        if (ERR_COMMON_SUCCESS != iErr)
        {
            iErr = ERR_COMMON_OPENSSL_FAILED;
            LOCKET_LOGERROR ("EVP_PKEY_decrypt failed. %s", LOCKET_ERR_GetString());
            break;
        }

        if (pucConvertOutput != pstRSA->pucOutput)
        {
            pRet = LOCKET_CIPHER_ByteToOther (pstRSA->iOutputEncode, pucConvertOutput, 
                                             iOutputLen, pstRSA->pucOutput, &pstRSA->iOutputLen);
            if (NULL == pRet)
            {
                iErr = ERR_COMMON_CONVERT_FAILED;
                break;
            }
        }
        else
        {
            pstRSA->iOutputLen = iOutputLen;
        }
    } while (0);

    /* don't forget cleanup */
    if (ENCODE_BINARY != pstRSA->iInputEncode)
    {
        LOCKET_free ((void**) &pucConvertInput);
    }
    if (ENCODE_BINARY != pstRSA->iOutputEncode)
    {
        LOCKET_free ((void**) &pucConvertOutput);
    }
    EVP_PKEY_CTX_free(pstCtx);
    EVP_PKEY_free(pstKey);

    return iErr;
}
